

/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"
#include <cassert>
#include <queue>
using std::queue;

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = -1;
	treeHeight = 0;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	//open the inde file
	RC rc = pf.open(indexname, mode);
	if (rc)
		return rc;

	//get auxiliary information
	char buffer[PageFile::PAGE_SIZE];
	rc = pf.read(0, buffer);
	if (rc == 0) {
		//get the root pid
		memcpy(&rootPid, buffer, sizeof(PageId));

		//get the tree hieght
		memcpy(&treeHeight, buffer + sizeof(PageId), sizeof(int));

		if ((treeHeight > 0 && (rootPid < 1 || rootPid >= pf.endPid()))
			|| treeHeight < 0) {
			pf.close();
			return RC_INVALID_FILE_FORMAT;
		}
	}
	else if (rc == RC_INVALID_PID) {
		//the index is empty.  Initialize it.
		rootPid = -1;
		treeHeight = 0;

		//write rootPid and treeHeight to disk
		if (rc = _updateAuxiliaryPage(rootPid, treeHeight)) {
			pf.close();
			return rc;
		}
	}
	else {
		//some other error occured
		pf.close();
		return rc;
	}
	
    return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
	RC rc;
	if (rc = pf.close()) {
		return rc;
	}
    return 0;
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	RC rc;
	//check that there is a root
	if (treeHeight > 0) {
		PageId siblingPid;
		int siblingKey;
		rc = _insert(key, rid, treeHeight-1, rootPid, siblingPid, siblingKey);
		if (rc != 0 && rc != RC_SPLIT_NODE)
			return rc;
		return 0;
	}
	else if (treeHeight == 0) {

		//The index is empty. Use a leaf node for the root
		BTLeafNode newRoot;
		if (rc = newRoot.insert(key, rid))
			return rc;

		//write the new root to disk
		PageId newRootPid = pf.endPid();
		if (rc = newRoot.write(newRootPid, pf))
			return rc;
		
		//update rootPid and treeHeight on disk
		if (rc = _updateAuxiliaryPage(newRootPid, treeHeight+1))
			return rc;
		treeHeight++;
		rootPid = newRootPid;
		
		return 0;
	}
	else {
		return RC_INVALID_FILE_FORMAT;
	}
	
	return 0;
}

RC BTreeIndex::_insert(int key, const RecordId& rid, int level, PageId nodePid, PageId& siblingPid, int& siblingKey) {

	RC rc;
	bool splitNode = false;
	if (level == 0) {
		//At a leaf node
		BTLeafNode leafNode;
		if (rc = leafNode.read(nodePid, pf))
			return rc;

		rc = leafNode.insert(key, rid);
		if (rc == 0) {
			//insert succeeded. write the node to disk
			leafNode.write(nodePid, pf);
		}
		if (rc == RC_NODE_FULL) {
			//try insert and split
			BTLeafNode sibling;
			if (rc = leafNode.insertAndSplit(key, rid, sibling, siblingKey))
				return rc;

			//link the old leaf node and its new sibling
			siblingPid = pf.endPid();
			sibling.setNextNodePtr(leafNode.getNextNodePtr());
			leafNode.setNextNodePtr(siblingPid);

			//write the nodes to disk
			sibling.write(siblingPid, pf);
			leafNode.write(nodePid, pf);

			//Check if the root node that was split.
			if (level == treeHeight-1) {
				//create a new root
				if (rc = _newRoot(nodePid, siblingKey, siblingPid))
					return rc; //at this point the index is corrupted
			}

			splitNode = true;
		}
		else if (rc) {
			//some other error occured
			return rc;
		}
	}
	else {
		//At a non-leaf node
		BTNonLeafNode node;
		if (rc = node.read(nodePid, pf))
			return rc;

		//Find the pointer to follow
		PageId childPid;
		if (rc = node.locateChildPtr(key, childPid))
			return rc;

		//Insert recursively into the child node
		PageId childSiblingPid;
		int childSiblingKey;
		rc = _insert(key, rid, level-1, childPid, childSiblingPid, childSiblingKey);

		if (rc == RC_SPLIT_NODE) {

			//The child node was split and generated a sibling node.  Insert
			//the sibling key and pid into this node.
			rc = node.insert(childSiblingKey, childSiblingPid);

			if (rc == 0) {
				//write the contents of this node to disk
				node.write(nodePid, pf);
			}
			else if (rc == RC_NODE_FULL) {
				//try insert and split
				BTNonLeafNode sibling;
				if (rc = node.insertAndSplit(childSiblingKey, childSiblingPid, sibling, siblingKey))
					return rc;

				//write the nodes to disk
				siblingPid = pf.endPid();
				sibling.write(siblingPid, pf);
				node.write(nodePid, pf);

				//Check if the root node that was split.
				if (level == treeHeight-1) {
					//create a new root
					if (rc = _newRoot(nodePid, siblingKey, siblingPid))
						return rc; //at this point the index is corrupted
				}

				splitNode = true;
			}
			else {
				//error occurred
				return rc;
			}
		}
		else if (rc) {
			//error occured
			return rc;
		}
	}

	return splitNode ? RC_SPLIT_NODE : 0;
}

RC BTreeIndex::_newRoot(PageId pid1, int key, PageId pid2) {
	
	BTNonLeafNode newRoot;
	RC rc;
	
	if (rc = newRoot.initializeRoot(pid1, key, pid2))
		return rc;
	
	PageId newRootPid = pf.endPid();
	if (rc = newRoot.write(newRootPid, pf))
		return rc;
	
	//update tree height and rootPid
	treeHeight++;
	rootPid = newRootPid;
	if (rc = _updateAuxiliaryPage(rootPid, treeHeight))
		return rc;

	return 0;
}

RC BTreeIndex::_updateAuxiliaryPage(PageId newRootPid, int newTreeHeight) {
	//write new treeHeight and rootPid to disk
	RC rc;
	char buffer[PageFile::PAGE_SIZE];
	memcpy(buffer, &newRootPid, sizeof(PageId));
	memcpy(buffer + sizeof(PageId), &newTreeHeight, sizeof(int));
	if (rc = pf.write(0, buffer)) {
		return rc;
	}
	return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	RC rc;
	if (rootPid < 1 || treeHeight < 1) {
		return RC_NO_SUCH_RECORD;
	}
	
	int pid = rootPid;
	if (treeHeight > 1) {
		BTNonLeafNode curNode;
		int height = treeHeight;
		while (height > 1) {
			if ((rc = curNode.read(pid, pf)) != 0)
				return rc;
			if ((rc = curNode.locateChildPtr(searchKey, pid)) != 0)
				return rc;
					
			height--;
		}
	}
	BTLeafNode leafNode;
	if ((rc = leafNode.read(pid, pf)) != 0)
		return rc;
	int eid;
	if ((rc = leafNode.locate(searchKey, eid)) != 0)
		return rc;
	cursor.eid = eid;
	cursor.pid = pid;
	
	return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	BTLeafNode leaf;
	RC rc;
	
	if (cursor.eid < 0 || cursor.pid < 0 || cursor.eid >= BT_MAX_NUM_KEYS)
		return RC_INVALID_CURSOR;
	if (cursor.pid == 0)
		return RC_END_OF_TREE;
	
	if ((rc = leaf.read(cursor.pid, pf)) != 0)
		return rc;
	if ((rc = leaf.readEntry(cursor.eid, key, rid)) != 0)
		return rc;
		
	if (cursor.eid == leaf.getKeyCount() - 1) {
		cursor.pid = leaf.getNextNodePtr();
		cursor.eid = 0;
	}
	else
		cursor.eid++;
	
    return 0;
}

struct Qval {
	PageId pid;
	int level;
};

void BTreeIndex::print() {
	queue<Qval> q;
	if (treeHeight == 0)
		return;
	Qval qval;
	qval.pid = rootPid;
	qval.level = treeHeight-1;
	q.push(qval);
	int level = -1;
	while (!q.empty()) {
		qval = q.front();
		if (qval.level != level) {
			//reached a new level
			printf("\n-- Level %d\n\n", qval.level);
			level = qval.level;
		}
		q.pop();

		printf("pid: %d    ", qval.pid);
		//print the node.
		if (qval.level > 0) {
			BTNonLeafNode node;
			node.read(qval.pid, pf);
			int eid = 0;
			int key;
			PageId pid;
			printf("(pid,key): ");
			while (node.readEntry(eid, key, pid) == 0) {
				printf("(%d,%d) ", pid, key);
				Qval qv;
				qv.pid = pid;
				qv.level = level-1;
				q.push(qv);
				eid++;
			}

			pid = node.getLastNodePtr();
			Qval qv;
			qv.pid = pid;
			qv.level = level-1;
			q.push(qv);
			printf("(%d)\n", pid);
		}
		else {
			BTLeafNode node;
			node.read(qval.pid, pf);
			int eid = 0;
			int key;
			RecordId rid;
			printf("(rid.sid,key) ");
			while (node.readEntry(eid, key, rid) == 0) {
				printf("(%d,%d) ", rid.sid, key);
				eid++;
			}
			PageId pid = node.getNextNodePtr();
			printf("(%d)\n", pid);
		}


	}
}

void BTreeIndex::test() {
	RC rc;
	RecordId rid;
	rid.pid = 0;
	rid.sid = 70;

	IndexCursor ic;
	ic.pid = -1;
	ic.eid = -1;

	BTLeafNode node;
	int key;

	//Insert into an empty index
	assert(this->open("test.idx", 'w') == 0);
	assert(insert(70, rid) == 0);
	assert(this->close() == 0);
	assert(this->open("test.idx", 'r') == 0);
	assert(treeHeight == 1 && rootPid == 1);
	this->close();

// Important: the following tests are designed for N = 3

	//Insert into an index with only one node (and therefore only one level)
	//causing it to split and generate a second level.
	open("test.idx", 'w');
	rid.sid = 50;
	assert(!insert(50,rid));
	rid.sid = 60;
	assert(!insert(60,rid));

	assert(!locate(50, ic));
	assert(ic.pid == 1 && ic.eid == 0);
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key ==50 && rid.sid == 50);
	assert(node.getNextNodePtr() == 2);

	assert(!locate(60, ic));
	assert(ic.pid == 1 && ic.eid == 1);
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key ==60 && rid.sid == 60);

	assert(!locate(70, ic));
	assert(ic.pid == 2 && ic.eid == 0);
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key ==70 && rid.sid == 70);

	//No overflow
	rid.sid = 80;
	assert(!insert(80, rid));
	assert(!locate(80, ic));
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key == 80 && rid.sid == 80);

	//Split at leaf node, no overflow in next level up
	rid.sid = 65;
	assert(!insert(65, rid));
	assert(treeHeight == 2);
	assert(!locate(65, ic));
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key == 65 && rid.sid == 65);

	//Split at leaf, overflow the root
	rid.sid = 40;
	assert(!insert(40, rid));
	assert(treeHeight == 3);
	assert(!locate(40, ic));
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key == 40 && rid.sid == 40);

	//print();

	//Overflow a leaf node - insert 75

	rid.sid = 75;
	assert(!insert(75, rid));
	assert(treeHeight == 3);
	assert(!locate(75, ic));
	node.read(ic.pid, pf);
	node.readEntry(ic.eid, key, rid);
	assert(key == 75 && rid.sid == 75);

	close();
	unlink("test.idx");
	
	//Insert 100 tuples into an empty index
	int keys[1000];
	assert(!open("test.idx", 'w'));
	for (int i = 0; i < 1000; i++) {
		keys[i] = 1 + rand() % 1000000;
		rid.sid = keys[i];
		assert(!insert(keys[i], rid));
	}
	assert(!close());

	//Locate each of the 100 tuples
	assert(!open("test.idx", 'r'));
	for (int i = 0; i < 1000; i++) {
		assert(!locate(keys[i], ic));
		node.read(ic.pid, pf);
		node.readEntry(ic.eid, key, rid);
		assert(key == keys[i] && rid.sid == keys[i]);
	}
	assert(!close());
	unlink("test.idx");

	//Open a file for reading that doesn't exist
	assert(this->open("test.idx", 'r') == RC_FILE_OPEN_FAILED);

	//Open a file for writing that doesn't exist
	assert(this->open("test.idx", 'w') == 0);
	assert(rootPid == -1 && treeHeight == 0);
	assert(this->close() == 0);

	//Open a file for reading that exists
	rootPid = -10;
	treeHeight = -10;
	assert(this->open("test.idx", 'r') == 0);
	assert(rootPid == -1 && treeHeight == 0);
	assert(this->close() == 0);

	unlink("test.idx");

	rid.pid = 12;
	rid.sid = 13;
/*	
	open("test.idx", 'w');
	for (int i = 1; i < 83; i++)
		insert(i, rid);
	close();
	
	open("test.idx", 'r');
	IndexCursor cursor;
	locate(1, cursor);
	assert(cursor.eid == 0);
	locate(2, cursor);
	assert(cursor.eid == 1);
	locate(3, cursor);
	assert(cursor.eid == 2);
	locate(82, cursor);
	assert(cursor.eid == 81);
	assert(locate(83, cursor) == RC_NO_SUCH_RECORD);
	cursor.pid = -1;
	cursor.eid = -1;
	close();
	
	open("test.idx", 'w');
	insert(83, rid);
	close();
	
	open("test.idx", 'r');	
	locate(83, cursor);
	assert(cursor.pid == 2 && cursor.eid == 40);
	locate(42 ,cursor);
	assert(cursor.pid == 1 && cursor.eid == 41);
	locate(43, cursor);
	assert(cursor.pid == 2 && cursor.eid == 0);
	locate(1, cursor);
	assert(cursor.pid == 1 && cursor.eid == 0);
	
	int key;
	readForward(cursor, key, rid);
	assert(key == 1 && cursor.pid == 1 && cursor.eid == 1);
	locate(42 ,cursor);
	readForward(cursor, key, rid);
	assert(key == 42 && cursor.pid == 2 && cursor.eid == 0);
	locate(83, cursor);
	rc = readForward(cursor, key, rid);
	assert(rc == RC_END_OF_TREE);
	
	
	close();
	
	unlink("test.idx");
	*/
}
